#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Sat Mar 30 20:18:43 2024

@author: lixiao
"""

import numpy as np
from scipy.spatial.transform import Rotation as sR
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D

from syntheticScene import (
    generateCube,
    generateCameraTrajectory,
    draw_camera,
    synthesizeImages
  )
from perpectiveGeometry import (
    computeRelativePose,
    decompose_camera,
    SE3_to_se3,
    se3_to_SE3
  )
from triangulation import linearTriangulation
from optimize import (
    reprojection_error,
    lieJacobian,
    lm,
    gn,
    simple_lm as lm2
  )

from utils import (
    vectorizeLieParameters,
    unvectorizeLieParameters
  )

#%%

def compute_relative_transform(R1, t1, R2, t2):
    """
    Compute the relative transformation between two coordinate frames.

    :param R1: 3x3 rotation matrix for the first frame
    :param t1: 3x1 translation vector for the first frame
    :param R2: 3x3 rotation matrix for the second frame
    :param t2: 3x1 translation vector for the second frame
    :return: relative_R, relative_t
    """
    relative_R = R1.T@R2  
    relative_t = R1.T @ (t2 - t1)
    # T1 = np.r_[np.c_[R1, t1], np.atleast_2d(np.array([0,0,0,1]))]
    # T2 = np.r_[np.c_[R2, t2], np.atleast_2d(np.array([0,0,0,1]))]
    # rT = np.linalg.inv(T1) @ T2
    return relative_R, relative_t

''' main '''

np.set_printoptions(suppress=True, precision=6)

f = 330
alpha_ccd = 1
cx = 240
cy = 320
K = np.array([[alpha_ccd*f, 0, cx],
             [0, alpha_ccd*f, cy],
             [0, 0, 1]])

world_points_gt = generateCube(4, 16)
num_views = 8
num_points = world_points_gt.shape[0]
Rs, ts = generateCameraTrajectory(num_views, 25)
cam_locs = ts.T
Ps = np.zeros((num_views, 3, 4))

fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')

ax.scatter(world_points_gt[:, 0], world_points_gt[:, 1], world_points_gt[:, 2], 'filled')
ax.set_xlabel('X')
ax.set_ylabel('Y')
ax.set_zlabel('Z')
ax.set_xlim(-26,26)
ax.set_ylim(-26,26)
ax.set_zlim(-26,26)
ax.set_title('The synthetic scene and cameras')

ex = []
for i in range(num_views):
    R_cam_W = Rs[i]
    # print(sR.from_matrix(R_cam_W).as_euler('xyz', degrees=True))
    ex.append(np.c_[R_cam_W, cam_locs[:, i]])
draw_camera(ax, ex)
plt.show()

images = synthesizeImages(world_points_gt, K, Rs, ts)

fig, axes = plt.subplots(2, 4, figsize=(12, 8))  
fig.subplots_adjust(hspace=0.4, wspace=0.3)  
for i in range(num_views):
    row = i // 4  
    col = i % 4  
    ax = axes[row, col]  
    ax.scatter(images[i][0, :], images[i][1, :], marker='o')
    ax.set_xlabel('x')
    ax.set_ylabel('y')
    ax.set_aspect('equal')
    ax.set_title('Image {}'.format(i+1))
plt.tight_layout()
plt.show()


R_cur = np.eye(3)
t_cur = np.zeros((3, 1))
Ps = np.zeros((num_views, 3, 4))
Ts = np.zeros((num_views, 3, 4))
X_init = None # 3d points
for i in range(1, num_views):
    # Compute the relative pose of the current view with respect to the first view
    R_21, t_21 = computeRelativePose(images[0], images[i], K)
    
    # If this is the second view, use this view to initialize the initial structure estimate as well, by triangulation.
    if i == 1:
        Ps[0] = np.c_[K, np.array([0,0,0])]
        Ps[1] = K @ np.c_[R_21, t_21]
        
        Ts[0] = np.c_[np.eye(3), np.array([0,0,0])]
        Ts[1] = np.c_[R_21, t_21]
        
        X_init = linearTriangulation(images[0], images[i], Ps[0], Ps[1])
        X_init = X_init / np.tile(X_init[3], (4, 1))
    else:
        # Update rotation and translation for current view
        R_cur = R_21
        t_cur = t_21
        
        # Rescale translation vector
        P_temp = K @ np.c_[R_21, t_21]
        X_temp = linearTriangulation(images[0], images[i], Ps[0], P_temp)
        X_temp = X_temp / np.tile(X_temp[3], (4, 1))
        
        numSamplesForScale = 25
        scaleFactors = np.zeros(numSamplesForScale)
        for j in range(numSamplesForScale):
            sampledPoints = np.random.choice(num_points, 2, replace=False)
            d1 = np.sqrt(np.sum((X_init[:3, sampledPoints[0]] - X_init[:3, sampledPoints[1]])**2))
            di = np.sqrt(np.sum((X_temp[:3, sampledPoints[0]] - X_temp[:3, sampledPoints[1]])**2))
            scaleFactor = d1 / di
            scaleFactors[j] = scaleFactor
            
        avgScaleFactor = np.mean(scaleFactors)
        t_cur = avgScaleFactor * t_cur
        
        # Construct the projection matrix using the estimated relative pose
        Ps[i] = K @ np.c_[R_cur, t_cur]
        Ts[i] = np.c_[R_cur, t_cur]


relative_Ts = np.zeros((num_views, 3, 4))
for i in range(num_views):
  r, t = compute_relative_transform(Rs[i], ts[i],Rs[0], ts[0])
  relative_Ts[i] = np.c_[r, t]


X_gt = X_init.copy()
Ps_gt = Ps.copy()
X_init_add_noise = np.s_[0:3, :]
noise_std = 0.3
noisy_submatrix = X_init[X_init_add_noise] +\
                np.random.normal(0, noise_std, X_init[X_init_add_noise].shape)

X_init[X_init_add_noise] = noisy_submatrix

Ps = Ps + np.random.normal(0, 0.001, Ps.shape)


para_views = []
para_points = X_init[0:3].T

for i in range(num_views):
  dK, R, Pc, _, _ = decompose_camera(Ps[i])
  T = np.r_[np.c_[R, -R@Pc], np.array([[0,0,0,1]])]
  ksi = SE3_to_se3(T)
  para_views.append(ksi)
para_views = np.array(para_views)

para_vec = vectorizeLieParameters(para_views, para_points)

M = para_views.shape[0]
N = para_points.shape[0]

def warp_fx(para_vec, x, K, M, N):
  Ps, X = unvectorizeLieParameters(para_vec, K, M, N)
  return reprojection_error(Ps, X, x)

err = warp_fx(para_vec, images, K, M, N)
err = np.linalg.norm(err)
print('err before ba: ', err)

# Ps_opt, X_opt, err_opt = gn(para_vec, images, K, M, N, warp_fx, lieJacobian)
Ps_lm, X_lm, err_lm = lm2(para_vec, images, K, M, N, warp_fx, lieJacobian)
print('err after lm: ', err_lm)
Ps_opt, X_opt, err_opt = gn(para_vec, images, K, M, N)
print('err after ba: ', err_opt)

fig = plt.figure(figsize=(10, 5))
ax1 = fig.add_subplot(131, projection='3d')
ax1.scatter(X_init[0], X_init[1], X_init[2], c='r', marker='o')
ax1.set_title('Noised 3d Points')
# 第二个子图
ax2 = fig.add_subplot(132, projection='3d')
ax2.scatter(X_opt[0], X_opt[1], X_opt[2], c='b', marker='o')
ax2.set_title('After GN-BA 3d Points')

ax3 = fig.add_subplot(133, projection='3d')
ax3.scatter(X_lm[0], X_lm[1], X_lm[2], c='b', marker='o')
ax3.set_title('After LM-BA 3d Points')

#%%
































